import os
import json
import shutil
from utils.img_db import ImgDB

# TODO:添加logging

DATA_SET = "imgs"
RESULT_DIR = os.path.join("data", DATA_SET, "result")
IMG_DIR = os.path.join("data", DATA_SET, "lmdb", "eval", "imgs")
IMG_FEAT_PATH = os.path.join("data", DATA_SET, "eval_imgs.img_feat.jsonl")
TXT_FEAT_PATH = os.path.join("data", DATA_SET, "query_texts.txt_feat.jsonl")
PRED_PATH = os.path.join("data", DATA_SET, "query_predictions.jsonl")
RESUME_PATH = "data/pretrained_weights/clip_cn_rn50.pt"

if not os.path.exists(RESULT_DIR):
    os.makedirs(RESULT_DIR)

from CLIP.feature_extractor import (
    ImageConfig,
    TextConfig,
    extract_text_features,
    extract_image_features,
)
from CLIP.topk_predictions import make_topk_predictions, TopKConfig


def run_retrieval(query: str, db: ImgDB, top_k: int = 10):
    temp_txt = os.path.join("data", DATA_SET, "query_temp.jsonl")
    with open(temp_txt, "w", encoding="utf-8") as f:
        f.write(
            json.dumps(
                {"text_id": 0, "text": query, "image_ids": []},
                ensure_ascii=False,
            )
            + "\n"
        )
    # 提取文本特征
    extract_text_features(
        TextConfig(
            text_data=temp_txt,
            output=TXT_FEAT_PATH,
            resume=RESUME_PATH,
            vision_model="RN50",
            text_model="RBT3-chinese",
            context_length=52,
            text_batch_size=1,
            gpu=0,
            debug=False,
        )
    )
    # 检索图片
    make_topk_predictions(
        TopKConfig(
            image_feats=IMG_FEAT_PATH,
            text_feats=TXT_FEAT_PATH,
            output=PRED_PATH,
            top_k=top_k,
            eval_batch_size=32768,
            debug=False,
        )
    )
    # 4. 读取预测结果
    # 清空 result 目录
    shutil.rmtree(RESULT_DIR)
    os.makedirs(RESULT_DIR)
    with open(PRED_PATH, "r", encoding="utf-8") as f:
        preds = [json.loads(line) for line in f]
    top_imgs = preds[0]["image_ids"] if preds else []
    # 5. 读取图片文件名列表
    img_files = db.list_images()
    # 6. 复制前 top_k 张图片到 result 目录
    for idx, img_id in enumerate(top_imgs[:top_k]):
        src = db.get_image_path(img_id)
        dst = os.path.join(RESULT_DIR, str(idx) + "." + src.rsplit(".", 1)[-1])
        shutil.copy2(src, dst)
    print(f"已将最接近的{top_k}张图片复制到 {RESULT_DIR}")


def main():
    db = ImgDB(
        root_path=os.path.join(os.path.dirname(__file__), "data", DATA_SET),
        img_feat_path=IMG_FEAT_PATH,
    )
    while True:
        query = input("请输入描述：").strip()
        run_retrieval(query, db)


if __name__ == "__main__":
    main()
